# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import argparse
import os
import pathlib
from typing import Generator, List, Optional, Tuple, cast

import gym.wrappers
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import torch

import mbrl
import mbrl.models
import mbrl.planning
import mbrl.util.common
from mbrl.models import ParallelGaussianMLP
from mbrl.env.ground_truth_mask import mask_dict

VisData = Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]


class Runner:
    def __init__(
            self,
            agent_dir: Optional[str],
    ):
        self.agent_dir = agent_dir

        self.cfg = mbrl.util.common.load_hydra_cfg(self.agent_dir)
        self.handler = mbrl.util.create_handler(self.cfg)

        self.env, term_fn, reward_fn = self.handler.make_env(self.cfg)

        agent_cfg = mbrl.util.common.load_hydra_cfg(agent_dir)
        self.agent = mbrl.planning.load_agent(agent_dir, self.env)

    def run(self):
        obs = self.env.reset()
        self.env.render()
        done = False
        total_reward = 0
        i = 0
        while not done:
            action = self.agent.act(obs)
            next_obs, reward, done, _ = self.env.step(action)
            self.env.render()
            total_reward += reward
            obs = next_obs
            i += 1

        print(total_reward)


if __name__ == "__main__":
    agent_dir = "../../exp/mbpo/time_scale-0.25/gym___Hopper-v2/2022.04.30/015718"

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--agent_dir",
        type=str,
        default=agent_dir,
        help="The directory where the agent configuration and data is stored. "
             "If not provided, a random agent will be used.",
    )
    args = parser.parse_args()

    runner = Runner(
        agent_dir=args.agent_dir,
    )
    runner.run()
